{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Preparation"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# sh data_load.sh \r\n",
    "!git clone https://github.com/aiyasin/LaneMLP.git"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "!cd LaneMLP/"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Merge Config"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "from configs.config import get_config\r\n",
    "config = get_config('./configs/yaml_cfg/lanemlp_24_224.yaml')"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "merging config from ./configs/yaml_cfg/lanemlp_24_224.yaml\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "config"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Data Test"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "from load_data.datasets_1 import get_dataset, get_dataloader\r\n",
    "dataset_train = get_dataset(config, mode='train')\r\n",
    "dataloader_train = get_dataloader(config, dataset_train, 'train', False)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "datas = []\r\n",
    "labels = []\r\n",
    "segs = []\r\n",
    "n = 0\r\n",
    "for data,label,seg in dataset_train:\r\n",
    "    datas.append(data)\r\n",
    "    labels.append(label)\r\n",
    "    segs.append(seg)\r\n",
    "    n=n+1\r\n",
    "    if n==1:\r\n",
    "        break\r\n",
    "print('data: ',datas[0].shape,\"\\n\", \"label: \",labels[0].shape,'\\n',segs[0].shape)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "n=0\r\n",
    "for batch_id, data in enumerate(dataloader_train):\r\n",
    "        image = data[0]\r\n",
    "        label = data[1]\r\n",
    "        print(data)\r\n",
    "        n=n+1\r\n",
    "        if n==1:\r\n",
    "                break"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Model"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "from model.lane_mlp import build_lane_mlp as build_model\r\n",
    "model = build_model(config)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "source": [
    "def inference(model, data_label, deploy):\r\n",
    "    if not deploy:\r\n",
    "        # print('train')\r\n",
    "        img, cls_label, seg_label = data_label[0], data_label[1], data_label[2]\r\n",
    "        img, cls_label, seg_label = img, cls_label.astype('int64'), seg_label.astype('int64')\r\n",
    "        cls_out, seg_out = model(img)\r\n",
    "        return {'cls_out': cls_out, 'cls_label': cls_label, 'seg_out':seg_out, 'seg_label': seg_label}\r\n",
    "    else:\r\n",
    "        img, cls_label = data_label[0], data_label[1]\r\n",
    "        img, cls_label = img, cls_label.astype('int64')\r\n",
    "        cls_out = model(img)\r\n",
    "        return {'cls_out': cls_out, 'cls_label': cls_label}"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Loss"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "import paddle\r\n",
    "import paddle.nn as nn\r\n",
    "from model.loss import SoftmaxFocalLoss, ParsingRelationLoss, ParsingRelationDis"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "def calc_loss(loss_dict, results):\r\n",
    "    loss = 0\r\n",
    "    for i in range(len(loss_dict['name'])):\r\n",
    "        data_src = loss_dict['data_src'][i]\r\n",
    "        datas = [results[src] for src in data_src]\r\n",
    "        # print(len(*datas))\r\n",
    "        loss_cur = loss_dict['op'][i](*datas)\r\n",
    "        \r\n",
    "        # if global_step % 20 == 0:\r\n",
    "        #     logger.add_scalar('loss/'+loss_dict['name'][i], loss_cur, global_step)\r\n",
    "        loss += loss_cur * loss_dict['weight'][i]\r\n",
    "    return loss"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "source": [
    "n=0\r\n",
    "for batch_id, data in enumerate(dataloader_train):\r\n",
    "        results = inference(model, data, deploy=False)\r\n",
    "        # loss = calc_loss(loss_dict, results)\r\n",
    "        # print(loss)\r\n",
    "        n=n+1\r\n",
    "        if n==1:\r\n",
    "                break"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "print(results['cls_out'].shape,results['cls_label'].shape,results['seg_out'].shape,results['seg_label'].shape)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "[1, 101, 56, 4] [1, 56, 4] [1, 36, 100] [1, 36, 100]\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "source": [
    "loss_dict = {\r\n",
    "            'name': ['cls_loss', 'relation_loss', 'relation_dis'],\r\n",
    "            'op': [SoftmaxFocalLoss(2), ParsingRelationLoss(), ParsingRelationDis()],\r\n",
    "            'weight': [1.0, 1.0, 0.0],\r\n",
    "            'data_src': [('cls_out', 'cls_label'), ('cls_out',), ('cls_out',)]\r\n",
    "        }"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "loss = calc_loss(loss_dict, results)\r\n",
    "print(loss)"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Tensor(shape=[1], dtype=float32, place=CUDAPlace(0), stop_gradient=False,\n",
      "       [5.26495886])\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Train"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "source": [
    "\r\n",
    "#   Copyright (c) 2021 PPViT Authors. All Rights Reserved.\r\n",
    "#\r\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n",
    "# you may not use this file except in compliance with the License.\r\n",
    "# You may obtain a copy of the License at\r\n",
    "#\r\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\r\n",
    "#\r\n",
    "# Unless required by applicable law or agreed to in writing, software\r\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n",
    "# See the License for the specific language governing permissions and\r\n",
    "# limitations under the License.\r\n",
    "\r\n",
    "\"\"\"ResMLP training/validation using single GPU \"\"\"\r\n",
    "\r\n",
    "import sys\r\n",
    "import os\r\n",
    "import time\r\n",
    "import logging\r\n",
    "import argparse\r\n",
    "import random\r\n",
    "import numpy as np\r\n",
    "import paddle\r\n",
    "import paddle.nn as nn\r\n",
    "import paddle.nn.functional as F\r\n",
    "from load_data.datasets_1 import get_dataloader\r\n",
    "from load_data.datasets_1 import get_dataset\r\n",
    "from model.lane_mlp import build_lane_mlp as build_model\r\n",
    "from utils.meter import AverageMeter\r\n",
    "from utils.meter import WarmupCosineScheduler\r\n",
    "from configs.config import get_config\r\n",
    "from configs.config import update_config\r\n",
    "from model.loss import get_loss_dict\r\n",
    "\r\n",
    "log_format = \"%(asctime)s %(message)s\"\r\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO,\r\n",
    "                    format=log_format, datefmt=\"%m%d %I:%M:%S %p\")\r\n",
    "\r\n",
    "config = get_config('./configs/yaml_cfg/lanemlp_24_224.yaml')\r\n",
    "\r\n",
    "if not os.path.exists(config.SAVE):\r\n",
    "    os.makedirs(config.SAVE, exist_ok=True)\r\n",
    "\r\n",
    "# set logging format\r\n",
    "logger = logging.getLogger()\r\n",
    "fh = logging.FileHandler(os.path.join(config.SAVE, 'log.txt'))\r\n",
    "fh.setFormatter(logging.Formatter(log_format))\r\n",
    "logger.addHandler(fh)\r\n",
    "logger.info(f'config= {config}')\r\n",
    "\r\n",
    "def inference(model, data_label, deploy):\r\n",
    "    if not deploy:\r\n",
    "        img, cls_label, seg_label = data_label[0], data_label[1], data_label[2]\r\n",
    "        img, cls_label, seg_label = img, cls_label.astype('int64'), seg_label.astype('int64')\r\n",
    "        cls_out, seg_out = model(img)\r\n",
    "        return {'cls_out': cls_out, 'cls_label': cls_label, 'seg_out':seg_out, 'seg_label': seg_label}\r\n",
    "    else:\r\n",
    "        img, cls_label = data_label[0], data_label[1]\r\n",
    "        img, cls_label = img, cls_label.astype('int64')\r\n",
    "        cls_out = model(img)\r\n",
    "        return {'cls_out': cls_out, 'cls_label': cls_label}\r\n",
    "\r\n",
    "def resolve_val_data(results):\r\n",
    "    results['cls_out'] = paddle.argmax(results['cls_out'], axis=1)\r\n",
    "    # if not deploy:\r\n",
    "    #     results['seg_out'] = paddle.argmax(results['seg_out'], axis=1)\r\n",
    "    return results[\"cls_out\"]\r\n",
    "\r\n",
    "def calc_loss(loss_dict, results):\r\n",
    "    loss = 0\r\n",
    "    for i in range(len(loss_dict['name'])):\r\n",
    "        data_src = loss_dict['data_src'][i]\r\n",
    "        datas = [results[src] for src in data_src]\r\n",
    "        # datas = datas.astype('int64')\r\n",
    "        loss_cur = loss_dict['op'][i](*datas)\r\n",
    "        # if global_step % 20 == 0:\r\n",
    "        #     logger.add_scalar('loss/'+loss_dict['name'][i], loss_cur, global_step)\r\n",
    "        loss += loss_cur * loss_dict['weight'][i]\r\n",
    "    return loss\r\n",
    "\r\n",
    "def train(dataloader,\r\n",
    "          model,\r\n",
    "          criterion,\r\n",
    "          loss_dict,\r\n",
    "          optimizer,\r\n",
    "          epoch,\r\n",
    "          total_batch,\r\n",
    "          debug_steps=100,\r\n",
    "          accum_iter=1,\r\n",
    "          ):\r\n",
    "    \"\"\"Training for one epoch\r\n",
    "    Args:\r\n",
    "        dataloader: paddle.io.DataLoader, dataloader instance\r\n",
    "        model: nn.Layer, a ViT model\r\n",
    "        criterion: nn.criterion\r\n",
    "        epoch: int, current epoch\r\n",
    "        total_epoch: int, total num of epoch, for logging\r\n",
    "        debug_steps: int, num of iters to log info\r\n",
    "        accum_iter: int, num of iters for accumulating gradients\r\n",
    "    Returns:\r\n",
    "        train_loss_meter.avg\r\n",
    "        train_acc_meter.avg\r\n",
    "        train_time\r\n",
    "    \"\"\"\r\n",
    "    model.train()\r\n",
    "    train_loss_meter = AverageMeter()\r\n",
    "    train_acc_meter = AverageMeter()\r\n",
    "    time_st = time.time()\r\n",
    "    \r\n",
    "    for batch_id, data in enumerate(dataloader):\r\n",
    "        image = data[0]\r\n",
    "        label = data[1]\r\n",
    "        results = inference(model, data, deploy=config.MODEL.DEPLOY)\r\n",
    "        # output = model(image)\r\n",
    "        loss = criterion(loss_dict, results)\r\n",
    "\r\n",
    "        #NOTE: division may be needed depending on the loss function\r\n",
    "        # Here no division is needed:\r\n",
    "        # default 'reduction' param in nn.CrossEntropyLoss is set to 'mean'\r\n",
    "        #loss =  loss / accum_iter\r\n",
    "\r\n",
    "        loss.backward()\r\n",
    "\r\n",
    "        if ((batch_id +1) % accum_iter == 0) or (batch_id + 1 == len(dataloader)):\r\n",
    "            optimizer.step()\r\n",
    "            optimizer.clear_grad()\r\n",
    "\r\n",
    "        # pred = F.softmax(output)\r\n",
    "        pred = resolve_val_data(results)\r\n",
    "        acc = paddle.metric.accuracy(pred, label.unsqueeze(1))\r\n",
    "\r\n",
    "        batch_size = image.shape[0]\r\n",
    "        train_loss_meter.update(loss.numpy()[0], batch_size)\r\n",
    "        train_acc_meter.update(acc.numpy()[0], batch_size)\r\n",
    "\r\n",
    "        if batch_id % debug_steps == 0:\r\n",
    "            logger.info(\r\n",
    "                f\"Epoch[{epoch:03d}/{config.TRAIN.NUM_EPOCHS:03d}], \" +\r\n",
    "                f\"Step[{batch_id:04d}/{total_batch:04d}], \" +\r\n",
    "                f\"Avg Loss: {train_loss_meter.avg:.4f}, \" +\r\n",
    "                f\"Avg Acc: {train_acc_meter.avg:.4f}\")\r\n",
    "\r\n",
    "    train_time = time.time() - time_st\r\n",
    "    return train_loss_meter.avg, train_acc_meter.avg, train_time\r\n",
    "\r\n",
    "\r\n",
    "def validate(dataloader, model, criterion, total_batch, debug_steps=10):\r\n",
    "    \"\"\"Validation for whole dataset\r\n",
    "    Args:\r\n",
    "        dataloader: paddle.io.DataLoader, dataloader instance\r\n",
    "        model: nn.Layer, a ViT model\r\n",
    "        criterion: nn.criterion\r\n",
    "        total_epoch: int, total num of epoch, for logging\r\n",
    "        debug_steps: int, num of iters to log info\r\n",
    "    Returns:\r\n",
    "        val_loss_meter.avg\r\n",
    "        val_acc1_meter.avg\r\n",
    "        val_acc5_meter.avg\r\n",
    "        val_time\r\n",
    "    \"\"\"\r\n",
    "    model.eval()\r\n",
    "    val_loss_meter = AverageMeter()\r\n",
    "    val_acc1_meter = AverageMeter()\r\n",
    "    val_acc5_meter = AverageMeter()\r\n",
    "    time_st = time.time()\r\n",
    "\r\n",
    "    with paddle.no_grad():\r\n",
    "        for batch_id, data in enumerate(dataloader):\r\n",
    "            image = data[0]\r\n",
    "            label = data[1]\r\n",
    "\r\n",
    "            output = model(image)\r\n",
    "            loss = criterion(output, label)\r\n",
    "\r\n",
    "            pred = F.softmax(output)\r\n",
    "            acc1 = paddle.metric.accuracy(pred, label.unsqueeze(1))\r\n",
    "            acc5 = paddle.metric.accuracy(pred, label.unsqueeze(1), k=5)\r\n",
    "\r\n",
    "            batch_size = image.shape[0]\r\n",
    "            val_loss_meter.update(loss.numpy()[0], batch_size)\r\n",
    "            val_acc1_meter.update(acc1.numpy()[0], batch_size)\r\n",
    "            val_acc5_meter.update(acc5.numpy()[0], batch_size)\r\n",
    "\r\n",
    "            if batch_id % debug_steps == 0:\r\n",
    "                logger.info(\r\n",
    "                    f\"Val Step[{batch_id:04d}/{total_batch:04d}], \" +\r\n",
    "                    f\"Avg Loss: {val_loss_meter.avg:.4f}, \" +\r\n",
    "                    f\"Avg Acc@1: {val_acc1_meter.avg:.4f}, \" +\r\n",
    "                    f\"Avg Acc@5: {val_acc5_meter.avg:.4f}\")\r\n",
    "\r\n",
    "    val_time = time.time() - time_st\r\n",
    "    return val_loss_meter.avg, val_acc1_meter.avg, val_acc5_meter.avg, val_time\r\n",
    "\r\n",
    "\r\n",
    "def main():\r\n",
    "    # 0. Preparation\r\n",
    "    last_epoch = config.TRAIN.LAST_EPOCH\r\n",
    "    seed = config.SEED\r\n",
    "    paddle.seed(seed)\r\n",
    "    np.random.seed(seed)\r\n",
    "    random.seed(seed)\r\n",
    "    #paddle.set_device('gpu:0')\r\n",
    "    # 1. Create model\r\n",
    "    model = build_model(config)\r\n",
    "    #model = paddle.DataParallel(model)\r\n",
    "    # 2. Create train and val dataloader\r\n",
    "    dataset_train = get_dataset(config, mode='train')\r\n",
    "    dataset_val = get_dataset(config, mode='val')\r\n",
    "    dataloader_train = get_dataloader(config, dataset_train, 'train', False)\r\n",
    "    dataloader_val = get_dataloader(config, dataset_val, 'val', False)\r\n",
    "    # 3. Define criterion\r\n",
    "    # criterion = nn.CrossEntropyLoss()\r\n",
    "    criterion = calc_loss\r\n",
    "    loss_dict = get_loss_dict(config)\r\n",
    "    # 4. Define lr_scheduler\r\n",
    "    scheduler = None\r\n",
    "    if config.TRAIN.LR_SCHEDULER.NAME == \"warmupcosine\":\r\n",
    "        scheduler = WarmupCosineScheduler(learning_rate=config.TRAIN.BASE_LR,\r\n",
    "                                          warmup_start_lr=config.TRAIN.WARMUP_START_LR,\r\n",
    "                                          start_lr=config.TRAIN.BASE_LR,\r\n",
    "                                          end_lr=config.TRAIN.END_LR,\r\n",
    "                                          warmup_epochs=config.TRAIN.WARMUP_EPOCHS,\r\n",
    "                                          total_epochs=config.TRAIN.NUM_EPOCHS,\r\n",
    "                                          last_epoch=config.TRAIN.LAST_EPOCH,\r\n",
    "                                          )\r\n",
    "    elif config.TRAIN.LR_SCHEDULER.NAME == \"cosine\":\r\n",
    "        scheduler = paddle.optimizer.lr.CosineAnnealingDecay(learning_rate=config.TRAIN.BASE_LR,\r\n",
    "                                                             T_max=config.TRAIN.NUM_EPOCHS,\r\n",
    "                                                             last_epoch=last_epoch)\r\n",
    "    elif config.scheduler == \"multi-step\":\r\n",
    "        milestones = [int(v.strip()) for v in config.TRAIN.LR_SCHEDULER.MILESTONES.split(\",\")]\r\n",
    "        scheduler = paddle.optimizer.lr.MultiStepDecay(learning_rate=config.TRAIN.BASE_LR,\r\n",
    "                                                       milestones=milestones,\r\n",
    "                                                       gamma=config.TRAIN.LR_SCHEDULER.DECAY_RATE,\r\n",
    "                                                       last_epoch=last_epoch)\r\n",
    "    else:\r\n",
    "        logging.fatal(f\"Unsupported Scheduler: {config.TRAIN.LR_SCHEDULER}.\")\r\n",
    "        raise NotImplementedError(f\"Unsupported Scheduler: {config.TRAIN.LR_SCHEDULER}.\")\r\n",
    "    # 5. Define optimizer\r\n",
    "    if config.TRAIN.OPTIMIZER.NAME == \"SGD\":\r\n",
    "        if config.TRAIN.GRAD_CLIP:\r\n",
    "            clip = paddle.nn.ClipGradByGlobalNorm(config.TRAIN.GRAD_CLIP)\r\n",
    "        else:\r\n",
    "            clip = None\r\n",
    "        optimizer = paddle.optimizer.Momentum(\r\n",
    "            parameters=model.parameters(),\r\n",
    "            learning_rate=scheduler if scheduler is not None else config.TRAIN.BASE_LR,\r\n",
    "            weight_decay=config.TRAIN.WEIGHT_DECAY,\r\n",
    "            momentum=config.TRAIN.OPTIMIZER.MOMENTUM,\r\n",
    "            grad_clip=clip)\r\n",
    "    elif config.TRAIN.OPTIMIZER.NAME == \"AdamW\":\r\n",
    "        if config.TRAIN.GRAD_CLIP:\r\n",
    "            clip = paddle.nn.ClipGradByGlobalNorm(config.TRAIN.GRAD_CLIP)\r\n",
    "        else:\r\n",
    "            clip = None\r\n",
    "        optimizer = paddle.optimizer.AdamW(\r\n",
    "            parameters=model.parameters(),\r\n",
    "            learning_rate=scheduler if scheduler is not None else config.TRAIN.BASE_LR,\r\n",
    "            weight_decay=config.TRAIN.WEIGHT_DECAY,\r\n",
    "            beta1=config.TRAIN.OPTIMIZER.BETAS[0],\r\n",
    "            beta2=config.TRAIN.OPTIMIZER.BETAS[1],\r\n",
    "            epsilon=config.TRAIN.OPTIMIZER.EPS,\r\n",
    "            grad_clip=clip)\r\n",
    "    else:\r\n",
    "        logging.fatal(f\"Unsupported Optimizer: {config.TRAIN.OPTIMIZER.NAME}.\")\r\n",
    "        raise NotImplementedError(f\"Unsupported Optimizer: {config.TRAIN.OPTIMIZER.NAME}.\")\r\n",
    "    # 6. Load pretrained model or load resume model and optimizer states\r\n",
    "    if config.MODEL.PRETRAINED:\r\n",
    "        assert os.path.isfile(config.MODEL.PRETRAINED + '.pdparams')\r\n",
    "        model_state = paddle.load(config.MODEL.PRETRAINED+'.pdparams')\r\n",
    "        model.set_dict(model_state)\r\n",
    "        logger.info(f\"----- Pretrained: Load model state from {config.MODEL.PRETRAINED}\")\r\n",
    "\r\n",
    "    if config.MODEL.RESUME:\r\n",
    "        assert os.path.isfile(config.MODEL.RESUME+'.pdparams') is True\r\n",
    "        assert os.path.isfile(config.MODEL.RESUME+'.pdopt') is True\r\n",
    "        model_state = paddle.load(config.MODEL.RESUME+'.pdparams')\r\n",
    "        model.set_dict(model_state)\r\n",
    "        opt_state = paddle.load(config.MODEL.RESUME+'.pdopt')\r\n",
    "        optimizer.set_state_dict(opt_state)\r\n",
    "        logger.info(\r\n",
    "            f\"----- Resume: Load model and optmizer from {config.MODEL.RESUME}\")\r\n",
    "    # 7. Validation\r\n",
    "    if config.EVAL:\r\n",
    "        logger.info('----- Start Validating')\r\n",
    "        val_loss, val_acc1, val_acc5, val_time = validate(\r\n",
    "            dataloader=dataloader_val,\r\n",
    "            model=model,\r\n",
    "            criterion=criterion,\r\n",
    "            total_batch=len(dataloader_val),\r\n",
    "            debug_steps=config.REPORT_FREQ)\r\n",
    "        logger.info(f\"Validation Loss: {val_loss:.4f}, \" +\r\n",
    "                    f\"Validation Acc@1: {val_acc1:.4f}, \" +\r\n",
    "                    f\"Validation Acc@5: {val_acc5:.4f}, \" +\r\n",
    "                    f\"time: {val_time:.2f}\")\r\n",
    "        return\r\n",
    "    # 8. Start training and validation\r\n",
    "    logging.info(f\"Start training from epoch {last_epoch+1}.\")\r\n",
    "    for epoch in range(last_epoch+1, config.TRAIN.NUM_EPOCHS+1):\r\n",
    "        # train\r\n",
    "        logging.info(f\"Now training epoch {epoch}. LR={optimizer.get_lr():.6f}\")\r\n",
    "        train_loss, train_acc, train_time = train(dataloader=dataloader_train,\r\n",
    "                                                  model=model,\r\n",
    "                                                  criterion=criterion,\r\n",
    "                                                  optimizer=optimizer,\r\n",
    "                                                  epoch=epoch,\r\n",
    "                                                  total_batch=len(dataloader_train),\r\n",
    "                                                  debug_steps=config.REPORT_FREQ,\r\n",
    "                                                  accum_iter=config.TRAIN.ACCUM_ITER,\r\n",
    "                                                  loss_dict = loss_dict,\r\n",
    "                                                  )\r\n",
    "        scheduler.step()\r\n",
    "        logger.info(f\"----- Epoch[{epoch:03d}/{config.TRAIN.NUM_EPOCHS:03d}], \" +\r\n",
    "                    f\"Train Loss: {train_loss:.4f}, \" +\r\n",
    "                    f\"Train Acc: {train_acc:.4f}, \" +\r\n",
    "                    f\"time: {train_time:.2f}\")\r\n",
    "        # validation\r\n",
    "        if epoch % config.VALIDATE_FREQ == 0 or epoch == config.TRAIN.NUM_EPOCHS:\r\n",
    "            logger.info(f'----- Validation after Epoch: {epoch}')\r\n",
    "            val_loss, val_acc1, val_acc5, val_time = validate(\r\n",
    "                dataloader=dataloader_val,\r\n",
    "                model=model,\r\n",
    "                criterion=criterion,\r\n",
    "                total_batch=len(dataloader_val),\r\n",
    "                debug_steps=config.REPORT_FREQ)\r\n",
    "            logger.info(f\"----- Epoch[{epoch:03d}/{config.TRAIN.NUM_EPOCHS:03d}], \" +\r\n",
    "                        f\"Validation Loss: {val_loss:.4f}, \" +\r\n",
    "                        f\"Validation Acc@1: {val_acc1:.4f}, \" +\r\n",
    "                        f\"Validation Acc@5: {val_acc5:.4f}, \" +\r\n",
    "                        f\"time: {val_time:.2f}\")\r\n",
    "        # model save\r\n",
    "        if epoch % config.SAVE_FREQ == 0 or epoch == config.TRAIN.NUM_EPOCHS:\r\n",
    "            model_path = os.path.join(\r\n",
    "                config.SAVE, f\"{config.MODEL.TYPE}-Epoch-{epoch}-Loss-{train_loss}\")\r\n",
    "            paddle.save(model.state_dict(), model_path + '.pdparams')\r\n",
    "            paddle.save(optimizer.state_dict(), model_path + '.pdopt')\r\n",
    "            logger.info(f\"----- Save model: {model_path}.pdparams\")\r\n",
    "            logger.info(f\"----- Save optim: {model_path}.pdopt\")\r\n",
    "\r\n",
    "\r\n",
    "# if __name__ == \"__main__\":\r\n",
    "#     main()\r\n"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "merging config from ./configs/yaml_cfg/lanemlp_24_224.yaml\n",
      "0922 10:55:20 AM config= BASE: ['']\n",
      "DATA:\n",
      "  BATCH_SIZE: 1\n",
      "  BATCH_SIZE_EVAL: 8\n",
      "  CROP_PCT: 0.875\n",
      "  DATASET: TuSimple\n",
      "  DATA_PATH: D:\\Files\\Data\\TuSimple\n",
      "  DROPOUT: 0.1\n",
      "  GRIDING_NUM: 100\n",
      "  IMAGE_SIZE: [560, 1000]\n",
      "  MEAN: [0.485, 0.456, 0.406]\n",
      "  NUM_CLASSES: 1000\n",
      "  NUM_LANES: 4\n",
      "  NUM_WORKERS: 0\n",
      "  SEG_SIZE: [36, 100]\n",
      "  STD: [0.229, 0.224, 0.225]\n",
      "EVAL: False\n",
      "LOCAL_RANK: 0\n",
      "MODEL:\n",
      "  CLS_DIM: [101, 56, 4]\n",
      "  DEPLOY: False\n",
      "  HIDDEN_SIZE: 56\n",
      "  NAME: lanemlp_24_224\n",
      "  NUM_LAYERS: 24\n",
      "  PATCH_SIZE: [10, 10]\n",
      "  PRETRAINED: None\n",
      "  RESUME: None\n",
      "  TYPE: ResMLP\n",
      "NGPUS: 1\n",
      "REPORT_FREQ: 50\n",
      "SAVE: ./output\n",
      "SAVE_FREQ: 20\n",
      "SEED: 0\n",
      "TAG: default\n",
      "TRAIN:\n",
      "  ACCUM_ITER: 2\n",
      "  BASE_LR: 0.001\n",
      "  END_LR: 1e-05\n",
      "  GRAD_CLIP: 1.0\n",
      "  LAST_EPOCH: 0\n",
      "  LOSS:\n",
      "    SHP_W: 0.0\n",
      "    SIM_W: 1.0\n",
      "  LR_SCHEDULER:\n",
      "    DECAY_EPOCHS: 30\n",
      "    DECAY_RATE: 0.1\n",
      "    MILESTONES: 30, 60, 90\n",
      "    NAME: warmupcosine\n",
      "  NUM_EPOCHS: 300\n",
      "  OPTIMIZER:\n",
      "    BETAS: (0.9, 0.999)\n",
      "    EPS: 1e-08\n",
      "    MOMENTUM: 0.9\n",
      "    NAME: AdamW\n",
      "  WARMUP_EPOCHS: 3\n",
      "  WARMUP_START_LR: 1e-06\n",
      "  WEIGHT_DECAY: 0.01\n",
      "VALIDATE_FREQ: 20\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "main()"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "orig_nbformat": 4,
  "language_info": {
   "name": "python",
   "version": "3.7.11",
   "mimetype": "text/x-python",
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "pygments_lexer": "ipython3",
   "nbconvert_exporter": "python",
   "file_extension": ".py"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.7.11 64-bit ('paddle': conda)"
  },
  "interpreter": {
   "hash": "fefd0da3125638cae2156a8bb3a3c81b468e2ea575d5a90cf241edbddc1363cd"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}